import paramiko
import re
import time
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('huawei_config.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('HuaweiAutoConfig')


class HuaweiSwitchConfigurator:
    def __init__(self, ip, username, password, port=22):
        self.ip = ip
        self.username = username
        self.password = password
        self.port = port
        self.client = None
        self.channel = None
        self.version = None
        self.model = None
        self.command_map = {}

    def connect(self):
        """建立SSH连接并获取设备信息"""
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.client.connect(
                hostname=self.ip,
                port=self.port,
                username=self.username,
                password=self.password,
                timeout=15,
                banner_timeout=20,
                look_for_keys=False,
                allow_agent=False
            )

            # 创建交互式会话
            self.channel = self.client.invoke_shell()
            time.sleep(1)

            # 获取设备信息
            self._get_device_info()
            self._detect_command_variations()

            logger.info(f"成功连接 {self.ip} ({self.model}, VRP {self.version})")
            return True
        except Exception as e:
            logger.error(f"连接失败: {str(e)}")
            return False

    def _get_device_info(self):
        """获取设备型号和版本信息"""
        self._send_command("display version")
        output = self._receive_output()

        # 解析版本信息
        version_match = re.search(r'VRP \(R\) software, Version (\S+)', output)
        self.version = version_match.group(1) if version_match else "Unknown"

        # 解析设备型号
        model_match = re.search(r'^.*?(S[0-9]+[A-Z]*)\s', output, re.MULTILINE)
        self.model = model_match.group(1) if model_match else "Unknown"

    def _detect_command_variations(self):
        """检测命令变体并创建映射"""
        self.command_map = {
            'sys': 'system-view',
            'quit': 'quit',
            'save': 'save'
        }

        # 检测密码策略命令变体
        self._send_command("display current-configuration | include password policy")
        output = self._receive_output()

        if "local-aaa-user password policy" in output:
            self.command_map['password_policy'] = "local-aaa-user password policy"
        elif "password policy" in output:
            self.command_map['password_policy'] = "password policy"
        else:
            self.command_map['password_policy'] = "local-aaa-user password policy"  # 默认

        # 检测RSA密钥命令变体
        self._send_command("rsa local-key-pair ?")
        output = self._receive_output()

        if "create" in output:
            self.command_map['rsa_create'] = "rsa local-key-pair create"
        elif "generate" in output:
            self.command_map['rsa_create'] = "rsa local-key-pair generate"
        else:
            self.command_map['rsa_create'] = "rsa local-key-pair create"  # 默认

    def _send_command(self, command, wait=0.5):
        """发送命令到交换机"""
        if not self.channel:
            return False

        final_command = self.command_map.get(command, command)
        self.channel.send(final_command + "\n")
        time.sleep(wait)
        return True

    def _receive_output(self, timeout=5):
        """接收命令输出"""
        if not self.channel:
            return ""

        output = ""
        start_time = time.time()

        while time.time() - start_time < timeout:
            if self.channel.recv_ready():
                chunk = self.channel.recv(65535).decode('utf-8', 'ignore')
                output += chunk
            else:
                time.sleep(0.1)


        return output

    def _enter_system_view(self):
        """进入系统视图"""
        self._send_command("sys")
        return "Error: Failed to enter system view" not in self._receive_output()

    def configure_remote_management(self):
        """配置远程管理功能"""
        if not self._enter_system_view():
            return False

        commands = [
            "aaa",
            f"{self.command_map['password_policy']} administrator",
            "password history record number 0",
            "quit",
            f"local-user {self.username} password irreversible-cipher {self.password}",
            f"local-user {self.username} privilege level 15",
            f"local-user {self.username} service-type ssh telnet",
            "quit",
            "stelnet server enable",
            "user-interface vty 0 4",
            "authentication-mode aaa",
            "protocol inbound all",
            "quit",
            self.command_map['rsa_create'],
            "y",  # 确认创建密钥
            f"ssh user {self.username} authentication-type password",
            f"ssh user {self.username} service-type stelnet",
            "vlan batch 2 to 4000",
            "interface vlanif 303",
            "ip address",  # 需要实际IP地址
            "quit",
            "interface vlanif 1009",
            "quit",
            "ip route-static 0.0.0.0 0.0.0.0 10.3.3.254",
            "interface range GigabitEthernet 0/0/1 to GigabitEthernet 0/0/22",
            "port link-type access",
            "port default vlan 1009",
            "quit",
            "interface range GigabitEthernet 0/0/23 to GigabitEthernet 0/0/24",
            "port link-type trunk",
            "port trunk allow-pass vlan all",
            "quit",
            "quit",
            "save"
        ]

        for cmd in commands:
            self._send_command(cmd)
            output = self._receive_output()

            # 处理保存确认
            if cmd == "save" and "[Y/N]" in output:
                self._send_command("y")
                time.sleep(1)

        return True

    def configure_snmp(self):
        """配置SNMP协议"""
        if not self._enter_system_view():
            return False

        commands = [
            "undo snmp-agent sys-info version v1 v2c",
            "snmp-agent protocol source-status all-interface",
            "snmp-agent usm-user v3 xxgluser authentication-mode sha",
            "xxgl1688",  # 认证密码
            "xxgl1688",  # 确认密码
            "11",  # 加密协议选择
            "snmp-agent usm-user v3 xxgluser privacy-mode aes256",
            "xxgl1688.",  # 加密密码
            "xxgl1688.",  # 确认密码
            "11",  # 加密协议选择
            "quit",
            "save"
        ]

        for cmd in commands:
            self._send_command(cmd)
            self._receive_output()

        return True

    def configure_dot1x(self):
        """配置802.1x认证"""
        if not self._enter_system_view():
            return False

        commands = [
            "vlan 6",
            "interface vlanif 6",
            "quit",
            "radius-server template sangfor_radius",
            "radius-server shared-key cipher Weihua@1234567",
            "radius-server authentication 10.1.1.16 1812 weight 80",
            "radius-server accounting 10.1.1.16 1813 weight 80",
            "radius-attribute nas-ip 10.3.2.242",
            "radius-server detect-server interval 30",
            "quit",
            "aaa",
            "authentication-scheme radius",
            "authentication-mode radius none",
            "y",  # 确认
            "quit",
            "accounting-scheme radius",
            "accounting-mode radius",
            "quit",
            "domain sangfor_802.1x",
            "authentication-scheme radius",
            "accounting-scheme radius",
            "radius-server sangfor_radius",
            "quit",
            "quit",
            "domain sangfor_802.1x",
            "free-rule-template name default_free_rule",
            "free-rule 0 destination ip 192.168.1.85 mask 255.255.255.255",
            "free-rule 1 destination ip 192.168.1.58 mask 255.255.255.255",
            "free-rule 2 destination ip 192.168.1.59 mask 255.255.255.255",
            "free-rule 3 destination ip 10.1.1.16 mask 255.255.255.255",
            "quit",
            "dot1x-access-profile name 802.1x",
            "dot1x timer client-timeout 30",
            "quit",
            "authentication-profile name sangfor_802.1x_vl205",
            "dot1x-access-profile 802.1x",
            "access-domain sangfor_802.1x force",
            "undo authentication handshake",
            "authentication event authen-server-down action authorize vlan 205",
            "authentication event authen-server-up action re-authen",
            "quit",
            "interface range GigabitEthernet 1/0/11 to GigabitEthernet 1/0/20",
            "authentication-profile sangfor_802.1x_vl205",
            "quit",
            "quit",
            "save"
        ]

        for cmd in commands:
            self._send_command(cmd)
            self._receive_output()

        return True

    def disconnect(self):
        """断开连接"""
        try:
            if self.channel:
                self.channel.close()
            if self.client:
                self.client.close()
            logger.info("连接已关闭")
        except:
            pass


def backup_config(ip, username, password):
    """备份交换机配置"""
    try:
        # 创建备份目录
        backup_dir = "config_backups"
        os.makedirs(backup_dir, exist_ok=True)

        # 连接设备
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(ip, username=username, password=password)

        # 获取配置
        stdin, stdout, stderr = client.exec_command("display current-configuration")
        config = stdout.read().decode('utf-8')

        # 保存配置
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{ip}_{timestamp}.cfg"
        filepath = os.path.join(backup_dir, filename)

        with open(filepath, 'w') as f:
            f.write(config)

        logger.info(f"配置已备份到: {filepath}")
        return True
    except Exception as e:
        logger.error(f"备份失败: {str(e)}")
        return False


def main():
    # 交换机信息
    SWITCH_IP = "10.3.2.201"
    USERNAME = "admin"
    PASSWORD = "Yunwei@1688"

    # 步骤1: 备份当前配置
    logger.info("开始备份当前配置...")
    if not backup_config(SWITCH_IP, USERNAME, PASSWORD):
        logger.error("备份失败，终止操作")
        return

    # 步骤2: 执行配置
    configurator = HuaweiSwitchConfigurator(SWITCH_IP, USERNAME, PASSWORD)
    if configurator.connect():
        logger.info("开始配置远程管理...")
        if configurator.configure_remote_management():
            logger.info("远程管理配置完成")
        else:
            logger.error("远程管理配置失败")

        logger.info("开始配置SNMP...")
        if configurator.configure_snmp():
            logger.info("SNMP配置完成")
        else:
            logger.error("SNMP配置失败")

        logger.info("开始配置802.1x认证...")
        if configurator.configure_dot1x():
            logger.info("802.1x认证配置完成")
        else:
            logger.error("802.1x认证配置失败")

        configurator.disconnect()


if __name__ == "__main__":
    main()